"""
内容：基于多层感知器 (MLP) 的 softmax 多分类
日期：2020年7月7日
作者：Howie
"""

from numpy import argmax
import pandas as pd
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

from keras.utils import to_categorical, plot_model
from keras.models import Sequential
from keras.layers import Dense, Dropout

LABELS = ['Iris-setosa', 'Iris-versicolor', 'Iris-virginica']

# 超参
N_CLASSES = 3
N_COMPONENT = 4
DROP_PROB = 0.5
DENSE_UNITS_1 = 64
DENSE_UNITS_2 = 64
BASE_LR = 0.01
DECAY = 1e-6
MOMENTUM = 0.9
EPOCHS = 500
BATCH_SIZE = 32


def iris_type(s):
    """
    # 列转换函数
    :return:
    """
    cls_label = {
        'Iris-setosa': 0,
        'Iris-versicolor': 1,
        'Iris-virginica': 2
    }
    return cls_label[s]


def load_iris():
    """
    # 加载鸢尾花数据集
    :return: 划分好的训练集和测试集
    """
    data = pd.read_csv(
        '../dataset/iris/Iris.csv',
        sep=',',
        encoding='utf-8',
        converters={
            'Species': iris_type}).drop('Id', axis=1)
    X = data.drop('Species', axis=1).values
    Y = to_categorical(data.loc[:, 'Species'].values)

    X_train, X_test, Y_train, Y_test = train_test_split(
        X, Y, test_size=0.2, shuffle=True)
    print(
        "train on {} samples, test on {} samples".format(
            X_train.shape[0],
            X_test.shape[0]))

    return X_train, X_test, Y_train, Y_test


def hist_plot(hist, model_name):
    """
    # 可视化训练过程
    :param hist: history对象
    :return:
    """
    print(hist.history)
    fig, loss_ax = plt.subplots()
    acc_ax = loss_ax.twinx()
    # 每个训练周期的训练误差与验证误差
    loss_ax.plot(hist.history['loss'], 'y', label='Loss')
    loss_ax.set_ylim([0.0, 0.5])
    # 每个训练周期的训练精度与验证精度
    acc_ax.plot(hist.history['accuracy'], 'b', label='Acc')
    acc_ax.set_ylim([0.8, 1.0])
    # 横轴与纵轴
    loss_ax.set_xlabel('epoch')
    loss_ax.set_ylabel('loss')
    acc_ax.set_ylabel('accuracy')
    # 标签
    loss_ax.legend(loc='upper left')
    acc_ax.legend(loc='lower left')
    # 标题
    plt.title(model_name)
    # 保存
    plt.savefig('./logs/History_Demo5-2_' + model_name + '.pdf')
    # 展示
    plt.show()


class Model:
    """
    模型类
    """

    def __init__(self):
        """
        搭建模型
        """
        self.model = Sequential()
        self.model.add(
            Dense(
                units=DENSE_UNITS_1,
                input_dim=N_COMPONENT,
                activation='relu'))
        self.model.add(Dropout(rate=DROP_PROB))
        self.model.add(Dense(units=DENSE_UNITS_2, activation='relu'))
        self.model.add(Dropout(rate=DROP_PROB))
        self.model.add(Dense(units=N_CLASSES, activation='softmax'))
        plot_model(
            self.model,
            to_file="./logs/Model_Demo5-2.pdf",
            show_shapes=True)

    def train(self, X_train, Y_train,):
        """
        训练模型
        :param X_train: 训练集样本
        :param Y_train: 训练集标签
        :return:
        """
        self.model.compile(
            loss='categorical_crossentropy',
            optimizer='rmsprop',
            metrics=['accuracy'])
        hist = self.model.fit(
            X_train,
            Y_train,
            epochs=EPOCHS,
            batch_size=BATCH_SIZE)
        hist_plot(hist, model_name='Multi-Layer Perceptron')

    def evaluate(self, X_test, Y_test):
        """
        评价模型
        :param X_test: 测试集样本
        :param Y_test: 测试集标签
        :return:
        """
        loss_and_metrics = self.model.evaluate(
            X_test, Y_test, batch_size=BATCH_SIZE)
        print("Evaluation loss and metrics: {}".format(loss_and_metrics))

    def predcit(self, samples, labels):
        """
        调用模型
        :param samples: 样本
        :param labels: 标签
        :return:
        """
        predictions = self.model.predict_classes(samples)
        for i, prediction in enumerate(predictions):
            print("({}) Ground truth: {} \t Prediction: {}".format(
                i + 1, LABELS[int(argmax(labels[i]))], LABELS[int(prediction)]))


def main():
    """
    主函数
    :return:
    """
    X_train, X_test, Y_train, Y_test = load_iris()
    model = Model()
    model.train(X_train, Y_train)
    model.evaluate(X_test, Y_test)
    model.predcit(X_test[:10], Y_test[:10])


if __name__ == '__main__':
    main()
